"""
LangGraph ReAct 工作流定义 - 使用StateGraph和Human-in-the-loop
"""

import json
import operator
from typing import TypedDict, Annotated, Sequence, List, Dict

from langchain_core.messages import BaseMessage, ToolMessage, HumanMessage, AIMessage
from loguru import logger
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

from ..schemas.base import Message
from . import agent  # 导入新的agent模块


# 1. 定义图的状态
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]


# 2. 构建图
graph_builder = StateGraph(AgentState)

# 添加节点
graph_builder.add_node("agent", agent.agent_node)
# 这两个节点现在只是作为中断的锚点，不执行任何操作
graph_builder.add_node("tools_for_approval", lambda state: {})
graph_builder.add_node("tools_for_auto_execution", lambda state: {})

# 定义入口点
graph_builder.set_entry_point("agent")

# 定义边
graph_builder.add_conditional_edges(
    "agent",
    agent.route_tools,
    {
        "tools_for_approval": "tools_for_approval",
        "tools_for_auto_execution": "tools_for_auto_execution",
        END: END,
    },
)
# 注意：在中断后，工作流的恢复将由前端触发，所以这里不需要边返回agent

# 3. 获取工具节点
tool_node = agent.get_tool_nodes()

# 4. 编译图
# 使用内存存储会话历史
memory = MemorySaver()
# 关键：我们在两个路径上都中断，以便向前端发送指令
agent_executor = graph_builder.compile(
    checkpointer=memory, 
    interrupt_before=["tools_for_approval", "tools_for_auto_execution"]
)


# 5. 定义工作流执行和恢复函数
async def execute_workflow(initial_messages: List[Message], session_id: str):
    """
    执行或继续ReAct工作流，直到下一个中断点或结束。
    """
    logger.info(f"开始/继续执行工作流: session_id={session_id}")
    config = {"configurable": {"thread_id": session_id}}

    # 转换消息格式，并处理历史中的 tool 消息
    lc_messages = []
    for msg in initial_messages:
        if msg.role == "user":
            lc_messages.append(HumanMessage(content=msg.content))
        elif msg.role == "assistant":
            lc_messages.append(AIMessage(content=msg.content))
        elif msg.role == "tool":
            # 忽略历史中的工具消息，因为它们的效果已反映在后续AI回复中
            logger.debug(f"忽略历史中的工具消息: {msg.content}")
            continue
    
    graph_input = {"messages": lc_messages}
        
    try:
        async for event in agent_executor.astream_events(graph_input, config=config, version="v1"):
            kind = event["event"]
            if kind == "on_chat_model_stream":
                chunk = event["data"].get("chunk")
                if chunk and chunk.content:
                    yield json.dumps({
                        "type": "chat_response",
                        "content": chunk.content
                    })

        # 流结束后，检查图的状态以确定下一步操作
        snapshot = agent_executor.get_state(config)
        next_nodes = snapshot.next
        if next_nodes:
            tool_calls = snapshot.values['messages'][-1].tool_calls
            
            if "tools_for_approval" in next_nodes:
                logger.info(f"[{session_id}] 工作流中断，等待工具批准: {tool_calls}")
                yield json.dumps({
                    "type": "tool_approval_request",
                    "payload": {"tool_calls": [{"name": t["name"], "args": t["args"], "id": t["id"]} for t in tool_calls]}
                })
            elif "tools_for_auto_execution" in next_nodes:
                logger.info(f"[{session_id}] 工作流中断，请求自动执行: {tool_calls}")
                yield json.dumps({
                    "type": "tool_auto_execution_request",
                    "payload": {"tool_calls": [{"name": t["name"], "args": t["args"], "id": t["id"]} for t in tool_calls]}
                })
            else: # 如果没有下一个节点，说明工作流已结束
                logger.info(f"[{session_id}] 工作流正常结束。")
                yield json.dumps({"type": "conversation_end"})
        else:
            # 如果没有下一个节点，说明工作流已结束
            logger.info(f"[{session_id}] 工作流正常结束。")
            yield json.dumps({"type": "conversation_end"})

    except Exception as e:
        logger.error(f"[{session_id}] 执行工作流时发生意外错误: {str(e)}", exc_info=True)
        yield json.dumps({
            "type": "error",
            "content": f"处理您的请求时发生错误: {str(e)}"
        })


async def resume_workflow(tool_outputs: List[Dict], session_id: str):
    """
    在前端执行工具后，使用工具输出恢复工作流。
    """
    logger.info(f"恢复工作流: session_id={session_id} with tool_outputs={tool_outputs}")
    config = {"configurable": {"thread_id": session_id}}
    
    tool_messages = []
    for output in tool_outputs:
        content = output.output
        try:
            output_data = json.loads(content)
            if output_data.get("error") == "操作被用户拒绝。":
                content = "用户已拒绝执行此工具。请告知用户此操作已被取消，并询问接下来需要做什么。"
        except (json.JSONDecodeError, TypeError):
            pass
        
        tool_messages.append(ToolMessage(content=content, tool_call_id=output.tool_call_id))

    # 直接使用带有工具消息的输入来继续工作流
    # LangGraph会从中断点自动恢复
    async for event in agent_executor.astream_events({"messages": tool_messages}, config=config, version="v1"):
        kind = event["event"]
        if kind == "on_chat_model_stream":
            chunk = event["data"].get("chunk")
            if chunk and chunk.content:
                yield json.dumps({
                    "type": "chat_response",
                    "content": chunk.content
                })
    
    # 再次检查是否需要新的批准
    snapshot = agent_executor.get_state(config)
    next_nodes = snapshot.next
    if next_nodes:
        if "tools_for_approval" in next_nodes:
            tool_calls = snapshot.values['messages'][-1].tool_calls
            logger.info(f"[{session_id}] 工作流再次中断，等待工具批准: {tool_calls}")
            yield json.dumps({
                "type": "tool_approval_request",
                "payload": {"tool_calls": [{"name": t["name"], "args": t["args"], "id": t["id"]} for t in tool_calls]}
            })
        elif "tools_for_auto_execution" in next_nodes:
            tool_calls = snapshot.values['messages'][-1].tool_calls
            logger.info(f"[{session_id}] 工作流再次中断，请求自动执行: {tool_calls}")
            yield json.dumps({
                "type": "tool_auto_execution_request",
                "payload": {"tool_calls": [{"name": t["name"], "args": t["args"], "id": t["id"]} for t in tool_calls]}
            })
        else: # 如果没有下一个节点，说明工作流已结束
            logger.info(f"[{session_id}] 工作流在恢复后正常结束。")
            yield json.dumps({"type": "conversation_end"})
    else:
        # 如果没有下一个节点，说明工作流已结束
        logger.info(f"[{session_id}] 工作流在恢复后正常结束。")
        yield json.dumps({"type": "conversation_end"})
